# coding:utf-8
# author: 湖南涉外经济学院信息与机电工程学院
# date: 2022.09.09

from flask import Flask, render_template, request, jsonify
from flask import Blueprint, render_template
import requests
import json
import base64
from io import BytesIO
from PIL import Image

index_ernie_vilg = Blueprint("ernie_vilg", __name__)


# 设置允许的文件格式

# 接收页面传递参数
@index_ernie_vilg.route('/ernie_vilg', methods=['POST', 'GET'])  # 添加路由
def ernie_vilg():
    if request.method == 'POST':
        try:
            ernie_vilg_text = request.form.get('describe')
            selected_style = request.form.get('selected_color')
            print(ernie_vilg_text, '\n=================\n', selected_style, '\n=================\n', request.form.get)
            file_list = ernie_vilg_gen(ernie_vilg_text, selected_style)
            # file_list = [f'static/images/target/{fff}' for fff in file_list]
            print(file_list)
            sourcefile0 = 'static/images/target/result_0.png'
            sourcefile1 = 'static/images/target/result_1.png'
            sourcefile2 = 'static/images/target/result_2.png'
            sourcefile3 = 'static/images/target/result_3.png'
            sourcefile4 = 'static/images/target/result_4.png'
            sourcefile5 = 'static/images/target/result_5.png'
            for filename in file_list:
                print(filename)
            return render_template('ernie_vilg.html', sourcefile1=sourcefile1, sourcefile2=sourcefile2,
                                   sourcefile3=sourcefile3, sourcefile4=sourcefile4, sourcefile5=sourcefile5,
                                   sourcefile0=sourcefile0, file_list=file_list)
        except Exception as e:
            print(e, "*" * 50)
            return render_template('ernie_vilg.html')
    return render_template('ernie_vilg.html')


# 生成
def ernie_vilg_gen(ernie_vilg_text, selected_style):
    # 发送HTTP请求
    print("结果: selected_style=", selected_style)
    data = {'text_prompts': ernie_vilg_text, 'style': selected_style}
    headers = {"Content-type": "application/json"}
    try:
        print("使用config+cpufing方案")
        url = "http://127.0.0.1:8869/predict/ernie_vilg"  # 使用GPU+CPU方案
        r = requests.post(url=url, headers=headers, data=json.dumps(data))
        print(r)
    except:#KeyError, VauleError
        print("使用config+cpufing方案失败,启用onlycpuconfig方案")
        url = "http://127.0.0.1:8869/predict/ernie_vilg" # 使用onlyCPU方案
        r = requests.post(url=url, headers=headers, data=json.dumps(data))
        print(r)
    print('###===###===###===###')
    print('data = ', data)
    print('###===###===###===###')
    print(r.json())
    print('###===###===###===###')
    # 获取返回结果
    file_list = []
    for i, result in enumerate(r.json()["results"]):
        image = Image.open(BytesIO(base64.b64decode(result)))
        image.save('static/images/target/result_{}.png'.format(i))
        file_list.append('static/images/target/result_{}.png'.format(i))
        print('result_{}.png'.format(i))
    return file_list
